package org.apache.spark.sql.execution.datasources.rest

import org.apache.spark.{InterruptibleIterator, Partition, SparkContext, TaskContext}
import org.apache.spark.rdd.RDD
import org.apache.spark.sql.Row
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.sources.Filter
import org.apache.spark.sql.types.StructType
import org.apache.spark.util.{CompletionIterator, NextIterator}

import scala.collection.mutable.ArrayBuffer

class RestRDD(sc: SparkContext,
              schema: StructType,
              requireColumns: Array[String],
              parts: Array[Partition],
              params: Map[String, String],
              filters: Array[Filter]) extends RDD[InternalRow](sc, Nil) {

  override def compute(split: Partition, context: TaskContext): Iterator[InternalRow] = {
    val inputMetrics = context.taskMetrics().inputMetrics
    val data = RestUtil.queryData1(prunedColumn(), inputMetrics).asInstanceOf[Iterator[InternalRow]]
    CompletionIterator[InternalRow, Iterator[InternalRow]](
      new InterruptibleIterator(context, data), () -> {
        println("==== execute resource close")
      })
  }

  def prunedColumn() = {
    val fieldsMap = Map(schema.fields.map(x => (x.name, x)): _*)
    new StructType(requireColumns.map(name => fieldsMap(name)))
  }


  /*override def compute(split: Partition, context: TaskContext): Iterator[Row] = new NextIterator[Row] {

    context.addTaskCompletionListener(context => closeIfNeeded())
    val inputMetrics = context.taskMetrics().inputMetrics

    val data: Iterator[Row] = RestUtil.queryData()

    override protected def getNext(): Row = {
      if (data.hasNext) {
        data.next()
      } else {
        finished = true
        null.asInstanceOf[Row]
      }
    }

    override protected def close(): Unit = {}
  }*/

  override protected def getPartitions: Array[Partition] = this.parts
}

/*object RestRDD {

  def scanTable(sc: SparkContext,
                schema: StructType,
                requiredColumns: Array[String],
                parts: Array[Partition],
                params: Map[String, String],
                filters: Array[Filter]): RDD[InternalRow] = {
    new RestRDD(sc, schema, requiredColumns, parts, params, filters)
  }
}*/
